/*
 * Copyright (c) 2020-2099 the original author or authors. All rights reserve.
 */
package org.tianyun.cloud.mybatis.plugins;

import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.datasource.DataSourceException;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.tianyun.cloud.mybatis.plugins.PageHelper.Context;
import org.tianyun.cloud.mybatis.plugins.sqlparser.SqlParser;
import org.tianyun.cloud.mybatis.plugins.sqlparser.SqlParserFactory;
import org.tianyun.cloud.mybatis.plugins.utils.AcessUtils;
import org.tianyun.cloud.mybatis.plugins.utils.MappedStatementUtils;
import org.tianyun.cloud.utils.CursorPageBuilder;
import org.tianyun.cloud.utils.PageBuilder;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Mybatis分页插件
 *
 * @auther ebert_chan
 */
//@formatter:off
@Intercepts({
	@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class }),
	@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class })
})
//@formatter:on
public class MybatisPagePlugin implements Interceptor {

    /*
     * @see org.apache.ibatis.plugin.Interceptor#intercept(org.apache.ibatis.plugin. Invocation)
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 上下文中无分页标识，跳过
        Context<?> context = PageHelper.get();
        if (context == null) {
            return invocation.proceed();
        }
        return context.getPageBuilder() != null ? autoStartPage(invocation, context.getPageBuilder()) : injectPageContextParamters(invocation, context.getCursorPageBuilder());
    }

    /**
     * 注入分页请求参数，mapper.xml中可以直接使用：#{page.offset}, #{page.pageSize}
     *
     * @param invocation
     * @param builder
     * @return
     * @throws Throwable
     * @auther ebert_chan
     */
    @SuppressWarnings("unchecked")
    private Object injectPageContextParamters(Invocation invocation, CursorPageBuilder<?> builder) throws Throwable {
        if (builder != null) {
            Object[] args = invocation.getArgs();
            Object parameterObject = args[1];

            ContextParamMap paramMap;
            if (parameterObject instanceof ContextParamMap) {
                paramMap = (ContextParamMap) parameterObject;
            } else {
                paramMap = new ContextParamMap();
                if (parameterObject != null) {
                    if (parameterObject instanceof Map) {
                        paramMap.putAll((Map<? extends String, ? extends Object>) parameterObject);
                    } else {
                        paramMap.put(ContextParamMap.DEFAULT_ARG_NAME, parameterObject);
                    }
                } else {
                    paramMap.put(ContextParamMap.DEFAULT_ARG_NAME, null);
                }
                args[1] = paramMap;
            }

            Map<String, Object> page = new HashMap<>(2);
            page.put("offset", builder.getOffset());
            page.put("pageSize", builder.getPageSize());
            paramMap.put(ContextParamMap.DEFAULT_PAGE_NAME, page);
        }
        return invocation.proceed();
    }

    /**
     * 自动分页，改写sql
     *
     * @param invocation
     * @param builder
     * @return
     * @throws Throwable
     * @auther ebert_chan
     */
    private Object autoStartPage(Invocation invocation, PageBuilder<?> builder) throws Throwable {
        Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameterObject = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler<?> resultHandler = (ResultHandler<?>) args[3];
        Executor executor = (Executor) invocation.getTarget();

        Object result = null;
        try {
            DataSource dataSource = mappedStatement.getConfiguration().getEnvironment().getDataSource();
            SqlParser sqlParser = SqlParserFactory.resolveDataSource(dataSource);
            if (sqlParser == null) {
                throw new DataSourceException("Mybatis分页插件无法自动解析当前使用的数据库类型");
            }

            Configuration configuration = mappedStatement.getConfiguration();
            // count statement id
            String statementId = mappedStatement.getId() + "$COUNT";
            // 从源mappedStatement复制新的count mappedStatement，并缓存起来
            MappedStatement countMappedStatement = MappedStatementUtils.getCopyMappedStatement(configuration, mappedStatement, statementId);
            // 获取总记录数
            Long totals = this.executeCountSql(sqlParser, executor, countMappedStatement, parameterObject, resultHandler);
            builder.setTotals(totals);

            // 转换为分页sql
            BoundSql boundSql = sqlParser.selectPageBoundSql(mappedStatement, parameterObject, builder.getPageIndex(), builder.getPageSize());
            // 创建缓存key
            CacheKey cacheKey = executor.createCacheKey(mappedStatement, parameterObject, rowBounds, boundSql);
            // 执行分页语句
            result = executor.query(mappedStatement, parameterObject, RowBounds.DEFAULT, resultHandler, cacheKey, boundSql);

        } finally {
            PageHelper.clear();
        }
        return result;
    }

    /**
     * 执行count查询语句
     *
     * @param sqlParser
     * @param executor
     * @param countMappedStatement
     * @param parameter
     * @param resultHandler
     * @return
     * @throws SQLException
     * @auther ebert_chan
     */
    private Long executeCountSql(SqlParser sqlParser, Executor executor, MappedStatement countMappedStatement, Object parameter, ResultHandler<?> resultHandler)
            throws SQLException {
        BoundSql boundSql = countMappedStatement.getBoundSql(parameter);
        // 创建缓存key
        CacheKey cacheKey = executor.createCacheKey(countMappedStatement, parameter, RowBounds.DEFAULT, boundSql);
        // 转换为count汇总sql
        String sql = sqlParser.selectCountSql(boundSql.getSql());
        // 重新设置sql
        AcessUtils.setSql(boundSql, sql);
        // 执行select count语句
        Object result = executor.query(countMappedStatement, parameter, RowBounds.DEFAULT, null, cacheKey, boundSql);
        // 获取总记录数
        Long count = (Long) ((List<?>) result).get(0);
        return count;
    }

}
